from template import *
from collect import collect


import sys

itime = 0
if len(sys.argv) == 2:
	t = int(sys.argv[1])
elif len(sys.argv) > 2:
	print("error: should have one argument, as time step")
else:
	t = itime


def vdf_maxwell_x(vdf_velocity_array, Tx, vx, species):
    vdf_array = np.zeros(vdf_velocity_array.shape[0])
    m_ov_2T = species["mass"] / (2.0 * Tx * const.e)
    for i in range(0, vdf_array.shape[0]):
        vdf_array[i] = math.sqrt(m_ov_2T/3.14) * math.exp( -m_ov_2T * (vdf_velocity_array[i] - vx) * (vdf_velocity_array[i] - vx) )
    return vdf_array




ref_path = "conduction_limit/data"
ref_config = json.load( open(ref_path+"/parameter.inp", "r") )
species_e = ref_config["Species"][0]
species_D1 = ref_config["Species"][1]
print(species_e["mass"], species_D1["mass"])

vdf_position = ref_config["Diagnostic"][0]["vedf_position"]
print("vdf_position: ", vdf_position)
vdf_index = []
for position in vdf_position:
    index = int(position / ref_config["cell_length"][0])
    vdf_index.append(index)
print("vdf_index: ", vdf_index)
cell_point_number = int(ref_config["sim_length"][0] / ref_config["cell_length"][0])
print("cell_point_number: ", cell_point_number)


vdf_velocity_array_e = collect("/Diagnostic/", "vdf_velocity_array", path = ref_path)[t, 0, 0, :]
vdf_velocity_array_D1 = collect("/Diagnostic/", "vdf_velocity_array", path = ref_path)[t, 0, 1, :]

print("size of x: ", vdf_velocity_array_e.shape)

#xmin = x.min()
#xmax = x.max()

list_fe_sheath = []
list_fe_mid = []
list_fe_source = []
list_fD1_sheath = []
list_fD1_mid = []
list_fD1_source = []

list_label = ["Rc = 0.5", "Rc = 0.6", "Rc = 0.7", "Rc = 0.8", "Rc = 0.9"]

list_ylabel = ["f_e(source)", "f_D1(source)", "f_e(mid)", "f_D1(mid)", "f_e(sheath)", "f_D1(sheath)"]

list_path = ["conduction_limit_re0.5/data", "conduction_limit_re0.6/data", "conduction_limit_re0.7/data", "conduction_limit/data", "conduction_limit_re0.9/data"]
for path_item in list_path:
    val = collect("/Diagnostic/", "vdf_parallel_to_B", path = path_item)[t, :, :, :]

    list_fe_sheath.append(val[0, 0, :])
    list_fe_mid.append(val[3, 0, :])
    list_fe_source.append(val[5, 0, :])

    list_fD1_sheath.append(val[0, 1, :])
    list_fD1_mid.append(val[3, 1, :])
    list_fD1_source.append(val[5, 1, :])

Te = collect("/Fields/", "T_global_e_avg", path = ref_path)[t, 0, 0, :]
Ve = collect("/Fields/", "Vparallel_global_e_avg", path = ref_path)[t, 0, 0, :]
TD1 = collect("/Fields/", "T_global_D1_avg", path = ref_path)[t, 0, 0, :]
VD1 = collect("/Fields/", "Vparallel_global_D1_avg", path = ref_path)[t, 0, 0, :]

list_fe_maxwell = []
list_fD1_maxwell = []

fe_maxwell = vdf_maxwell_x(vdf_velocity_array_e, Te[vdf_index[0]], Ve[vdf_index[0]], species_e)
list_fe_maxwell.append(fe_maxwell)
fe_maxwell = vdf_maxwell_x(vdf_velocity_array_e, Te[vdf_index[3]], Ve[vdf_index[3]], species_e)
list_fe_maxwell.append(fe_maxwell)
fe_maxwell = vdf_maxwell_x(vdf_velocity_array_e, Te[vdf_index[5]], Ve[vdf_index[5]], species_e)
list_fe_maxwell.append(fe_maxwell)

fD1_maxwell = vdf_maxwell_x(vdf_velocity_array_D1, TD1[vdf_index[0]], VD1[vdf_index[0]], species_D1)
list_fD1_maxwell.append(fD1_maxwell)
fD1_maxwell = vdf_maxwell_x(vdf_velocity_array_D1, TD1[vdf_index[3]], VD1[vdf_index[3]], species_D1)
list_fD1_maxwell.append(fD1_maxwell)
fD1_maxwell = vdf_maxwell_x(vdf_velocity_array_D1, TD1[vdf_index[5]], VD1[vdf_index[5]], species_D1)
list_fD1_maxwell.append(fD1_maxwell)

##inite the fig of matplotlib
fig=plt.figure(figsize=(15,10))
fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.35)

##============ fe source ======================================================
ax=fig.add_subplot(3,2,1)

for i in range(len(list_fe_source)):
    line0=ax.plot(vdf_velocity_array_e, list_fe_source[i], label = list_label[i])

line0=ax.plot(vdf_velocity_array_e, list_fe_maxwell[2], label = list_label[i], linestyle = linestyles[1])

#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[0], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 

num1 = 0
num2 = 1.4
num3 = 2
num4 = 0
handles, labels = ax.get_legend_handles_labels()
ax.legend(flip(handles, 3), flip(labels, 3), framealpha = 0.2, bbox_to_anchor=(num1, num2), loc=num3, borderaxespad=num4, ncol = 3)


ax.annotate("(a)", xy=get_axis_limits(ax, x_scale = -0.1, y_scale = -0.15), annotation_clip=False)

##============ fD1 source ======================================================
ax=fig.add_subplot(3,2,2)

for i in range(len(list_fD1_source)):
    line0=ax.plot(vdf_velocity_array_D1, list_fD1_source[i], label = list_label[i])

line0=ax.plot(vdf_velocity_array_D1, list_fD1_maxwell[2], label = list_label[i], linestyle = linestyles[1])

#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[1], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.annotate("(b)", xy=get_axis_limits(ax, x_scale = -0.1, y_scale = -0.15), annotation_clip=False)

##============ fe mid ======================================================
ax=fig.add_subplot(3,2,3)

for i in range(len(list_fe_mid)):
    line0=ax.plot(vdf_velocity_array_e, list_fe_mid[i], label = list_label[i])

line0=ax.plot(vdf_velocity_array_e, list_fe_maxwell[1], label = list_label[i], linestyle = linestyles[1])

#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[2], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.annotate("(c)", xy=get_axis_limits(ax, x_scale = -0.1, y_scale = -0.15), annotation_clip=False)



##============ fD1 mid ======================================================
ax=fig.add_subplot(3,2,4)

for i in range(len(list_fD1_mid)):
    line0=ax.plot(vdf_velocity_array_D1, list_fD1_mid[i], label = list_label[i])

line0=ax.plot(vdf_velocity_array_D1, list_fD1_maxwell[1], label = list_label[i], linestyle = linestyles[1])

#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[3], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.annotate("(d)", xy=get_axis_limits(ax, x_scale = -0.1, y_scale = -0.15), annotation_clip=False)


##============ fe sheath ======================================================
ax=fig.add_subplot(3,2,5)

for i in range(len(list_fe_sheath)):
    line0=ax.plot(vdf_velocity_array_e, list_fe_sheath[i], label = list_label[i])

line0=ax.plot(vdf_velocity_array_e, list_fe_maxwell[0], label = list_label[i], linestyle = linestyles[1])

#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[4], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.annotate("(e)", xy=get_axis_limits(ax, x_scale = -0.1, y_scale = -0.15), annotation_clip=False)


##============ fD1 sheath ======================================================
ax=fig.add_subplot(3,2,6)

for i in range(len(list_fD1_sheath)):
    line0=ax.plot(vdf_velocity_array_D1, list_fD1_sheath[i], label = list_label[i])

line0=ax.plot(vdf_velocity_array_D1, list_fD1_maxwell[0], label = list_label[i], linestyle = linestyles[1])

#ax.set_ylim(bottom = 0.0)
ax.set_ylabel(list_ylabel[5], fontsize = label_fontsize)
#ax.grid(True)
ax.set_yscale('log') 
#ax.legend(framealpha = 0.2)
ax.annotate("(f)", xy=get_axis_limits(ax, x_scale = -0.1, y_scale = -0.15), annotation_clip=False)


fig_file_name = "Vdf_all_cl_" + str(t) + ".svg"
fig.savefig(fig_file_name, dpi = 300)
##fig.show()       #when the program finishes,the figure disappears
#plt.axis('equal')
#plt.show()         #The command is OK




